{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VAE Training - Faces dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from glob import glob\n",
    "import numpy as np\n",
    "\n",
    "from models.VAE import VariationalAutoencoder\n",
    "from keras.preprocessing.image import ImageDataGenerator\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run params\n",
    "section = 'vae'\n",
    "run_id = '0001'\n",
    "data_name = 'faces'\n",
    "RUN_FOLDER = 'run/{}/'.format(section)\n",
    "RUN_FOLDER += '_'.join([run_id, data_name])\n",
    "\n",
    "if not os.path.exists(RUN_FOLDER):\n",
    "    os.mkdir(RUN_FOLDER)\n",
    "    os.mkdir(os.path.join(RUN_FOLDER, 'viz'))\n",
    "    os.mkdir(os.path.join(RUN_FOLDER, 'images'))\n",
    "    os.mkdir(os.path.join(RUN_FOLDER, 'weights'))\n",
    "\n",
    "mode =  'build' #'load' #\n",
    "\n",
    "\n",
    "DATA_FOLDER = './data/celeb/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "INPUT_DIM = (128,128,3)\n",
    "BATCH_SIZE = 32\n",
    "\n",
    "filenames = np.array(glob(os.path.join(DATA_FOLDER, '*/*.jpg')))\n",
    "\n",
    "NUM_IMAGES = len(filenames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_gen = ImageDataGenerator(rescale=1./255)\n",
    "\n",
    "data_flow = data_gen.flow_from_directory(DATA_FOLDER\n",
    "                                         , target_size = INPUT_DIM[:2]\n",
    "                                         , batch_size = BATCH_SIZE\n",
    "                                         , shuffle = True\n",
    "                                         , class_mode = 'input'\n",
    "                                         , subset = \"training\"\n",
    "                                            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae = VariationalAutoencoder(\n",
    "                input_dim = INPUT_DIM\n",
    "                , encoder_conv_filters=[32,64,64, 64]\n",
    "                , encoder_conv_kernel_size=[3,3,3,3]\n",
    "                , encoder_conv_strides=[2,2,2,2]\n",
    "                , decoder_conv_t_filters=[64,64,32,3]\n",
    "                , decoder_conv_t_kernel_size=[3,3,3,3]\n",
    "                , decoder_conv_t_strides=[2,2,2,2]\n",
    "                , z_dim=200\n",
    "                , use_batch_norm=True\n",
    "                , use_dropout=True)\n",
    "\n",
    "if mode == 'build':\n",
    "    vae.save(RUN_FOLDER)\n",
    "else:\n",
    "    vae.load_weights(os.path.join(RUN_FOLDER, 'weights/weights.h5'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae.encoder.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae.decoder.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LEARNING_RATE = 0.0005\n",
    "R_LOSS_FACTOR = 10000\n",
    "EPOCHS = 200\n",
    "PRINT_EVERY_N_BATCHES = 100\n",
    "INITIAL_EPOCH = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae.compile(LEARNING_RATE, R_LOSS_FACTOR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae.train_with_generator(     \n",
    "    data_flow\n",
    "    , epochs = EPOCHS\n",
    "    , steps_per_epoch = NUM_IMAGES / BATCH_SIZE\n",
    "    , run_folder = RUN_FOLDER\n",
    "    , print_every_n_batches = PRINT_EVERY_N_BATCHES\n",
    "    , initial_epoch = INITIAL_EPOCH\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gdl_code",
   "language": "python",
   "name": "gdl_code"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
